"""
以CIFAR10加载数据集的simple图片分类器
"""
import numpy as np
import torch
from torch import nn, optim
import torchvision
from torchvision import transforms
import matplotlib.pyplot as plt


class Net(nn.Module):
    """
    模板：神经网络构建模板
    """
    def __init__(self):
        super(Net, self).__init__()

        """
        https://pytorch-cn.readthedocs.io/zh/latest/package_references/torch-nn/#_1
        Conv2d:
        in_channels:  输入信号的通道
        out_channels: 卷积产生的通道
        kernel_size: 卷积核的尺寸
        stride=1: 卷积步长
        padding=0: 输入的每一条边补充0的层数
        dilation=1: 卷积核元素之间的间距
        groups=1: 从输入通道到输出通道的阻塞连接数
        
        input_shape: (num, channel, height, width)
        output_shale: 同上
        
        model的权重&偏置
        weight: 卷积的权重，大小是(out_channels, in_channels, kernel_size)
        bias: 卷积的偏置系数，大小是（out_channel）
        
        
        MaxPool2d: 池化
        kernel_size: max pooling的窗口大小
        stride=None: max pooling的窗口移动的步长。默认值是kernel_size
        padding=0:  输入的每一条边补充0的层数
        dilation=1: 一个控制窗口中元素步幅的参数
        return_indices=False: 如果等于True，会返回输出最大值的序号，对于上采样操作会有帮助
        ceil_mode=False: 如果等于True，计算输出信号大小的时候，会使用向上取整，代替默认的向下取整的操作
        
        input_shape: (num, channel, height, width)
        output_shale: 同上
        
        
        nn.Linear:对输入数据做线性变换
        in_features: 每个输入样本的大小
        out_features: 每个输出样本的大小
        bias: 若设置为False，这层不会学习偏置。默认值：True
        
        input_shape: (num, in_features)
        output_shape: (num, out_features)
        """
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=6, kernel_size=5)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5)

        # in_features = nn out channels * kernel_size
        self.fc1 = nn.Linear(in_features=16 * 5 * 5, out_features=120)
        self.fc2 = nn.Linear(in_features=120, out_features=84)
        self.fc3 = nn.Linear(in_features=84, out_features=10)

    def forward(self, x):
        """
        激活函数: nn.functional.relu
        """
        x = self.pool(nn.functional.relu(input=self.conv1(x)))
        x = self.pool(nn.functional.relu(input=self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)   # flatten
        x = nn.functional.relu(input=self.fc1(x))
        x = nn.functional.relu(input=self.fc2(x))
        x = self.fc3(x)

        return x


def load_data(path, is_download, num_workers):
    """
    数据加载
    :param path: 存储路径
    :param is_download: 是否下载
    :param num_workers: 进程数
    :return:
    """
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])
    train_data = torchvision.datasets.CIFAR10(root=path,
                                              train=True,
                                              download=is_download,
                                              transform=transform)
    train_loader = torch.utils.data.DataLoader(train_data,
                                               batch_size=4,
                                               shuffle=True,
                                               num_workers=num_workers)

    test_data = torchvision.datasets.CIFAR10(root=path,
                                             train=False,
                                             download=is_download,
                                             transform=transform)
    test_loader = torch.utils.data.DataLoader(test_data,
                                              batch_size=4,
                                              shuffle=False,
                                              num_workers=num_workers)
    return train_loader, test_loader


def img_show(img):
    """
    图片展示
    :param img:
    :return:
    """
    img = img / 2 + 0.5
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()


def display_test(data, classes):
    """
    展示图片与label对应关系
    :param data:
    :param classes:  目标类
    :return:
    """
    train_data_iter = iter(data)
    images, labels = train_data_iter.__next__()

    # display the image and correspond label
    labels_info = [classes[num] for num in labels]
    print(' '.join(labels_info))
    img_show(img=torchvision.utils.make_grid(images))


# train
def model_train(net, device, train_data, lr=0.001, momentum=0.9, epochs=2,
                cycle_nums=2000):
    """
    模型训练
    :param net:
    :param device:
    :param train_data:
    :param lr:
    :param momentum:
    :param epochs:
    :param cycle_nums:
    :return:
    """

    # 损失函数：交叉熵损失
    criterion = nn.CrossEntropyLoss()
    # 优化器
    optimizer = optim.SGD(net.parameters(), lr=lr, momentum=momentum)

    for epoch in range(epochs):
        running_loss = 0
        for idx, data in enumerate(train_data, 0):
            input, labels = data
            input, labels = input.to(device), labels.to(device)

            # 清零
            optimizer.zero_grad()

            outputs = net(input)
            loss = criterion(outputs, labels)
            loss.backward()
            # update weight
            optimizer.step()

            running_loss += loss.item()
            if idx % cycle_nums == (cycle_nums - 1):
                print('[{}, {}] loss: {}'.format(epoch + 1,
                                                 idx + 1,
                                                 running_loss / cycle_nums))
                running_loss = 0
    return net


def predict(net, images, classes):
    """
    预测
    :param net: 神经网络
    :param images:
    :param classes:
    :return:  预测的正确个数；预测的结果
    """
    # 预测每个图片在10个class下每个对应class的概率
    outputs = net(images)
    # 取概率最大的值对应的class为预测类
    _, pred = torch.max(outputs, 1)
    pred_info = [classes[num] for num in pred]
    return pred, ' '.join(pred_info)


def cifar10_run(device, path='../data/', is_download=True,
                num_workers=0, threds_num=4):
    """
    入口
    :param device: 程序运行环境
    :param path: 数据存储路径
    :param is_download:  是否下载
    :param num_workers:  加载数据时的线程数
    :param threds_num:  训练时的进程数
    :return:
    """
    train_loaders, test_loaders = load_data(path=path,
                                            is_download=is_download,
                                            num_workers=num_workers)
    classes = ('plan', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse',
               'ship', 'truck')

    # display_test(data=train_loaders, classes=classes)

    # 多线程
    torch.set_num_threads(threds_num)

    # 神经网络
    net = Net()
    net = net.to(device)

    # train
    net = model_train(net, device, train_data=train_loaders)

    """
    # test
    data_test_iter = iter(test_loaders)
    images, labels = data_test_iter.__next__()
    # predict
    pred, info = predict(net=net, images=images, classes=classes)
    jud = (pred == labels).sum().item()
    print(jud)
    print(info)
    
    labels_info = [classes[num] for num in labels]
    print('true label: {}'.format(' '.join(labels_info)))
    """

    # 准确率
    # 整体
    correct = 0
    total = 0
    # 各个分类
    class_correct = list(0. for _ in range(len(classes)))
    class_total = list(0. for _ in range(len(classes)))

    with torch.no_grad():
        for data in test_loaders:
            images, labels = data
            images, labels = images.to(device), labels.to(device)

            # pred
            pred, info = predict(net=net, images=images, classes=classes)
            total += labels.size(0)
            jud_num = (pred == labels).sum().item()
            correct += jud_num

            for idx, label in enumerate(labels):
                class_total[label] += 1
                if label == pred[idx]:
                    class_correct[label] += 1
    infos = 'total accuracy: {}'.format(100 * correct / total)
    info = ''
    for i in range(len(classes)):
        info += '\nAccuracy of {}: {}'.\
            format(classes[i], 100 * class_correct[i] / class_total[i])

    print('{}\n{}'.format(infos, info))
